package main

import (
	"fmt"
	"gopkg.in/yaml.v2"
	"io"
	"log"
	"log/slog"
	"net"
	"os"
)

const filePath = "./proxy.yaml"

type Proxy struct {
	Listen struct {
		Port int    `yaml:"port"`
		Host string `yaml:"host"`
	} `yaml:"listen"`
	Target struct {
		Port int    `yaml:"port"`
		Host string `yaml:"host"`
	} `yaml:"target"`
	ReadSize int `yaml:"readSize"`
}

var proxyMap = make(map[string]*Proxy)

func init() {
	file, err := os.ReadFile(filePath)
	if err != nil {
		log.Fatalf("读取%s文件失败：%v\n", filePath, err)
	}
	// 解析YAML文件
	err = yaml.Unmarshal(file, &proxyMap)
	if err != nil {
		log.Fatalf("解析%s文件失败：%v\n", filePath, err)
	}

}
func main() {
	for proxyName, proxy := range proxyMap {
		if proxyName == "" {
			log.Fatalln("代理名不能为空，请检查配置文件！")
		}
		if proxy.Target.Port == 0 || proxy.Target.Host == "" {
			log.Fatalf("【%s】 中代理Host和Port不能为空！\n", proxyName)
		}
		if proxy.Listen.Port == 0 || proxy.Listen.Host == "" {
			log.Fatalf("【%s】中监听Host和Port不能为空！\n", proxyName)
		}
		go ListenAndServe(proxyName, proxy)
	}
	select {}
}
func ListenAndServe(proxyName string, proxy *Proxy) {
	endpoint := fmt.Sprintf("%s:%d", proxy.Listen.Host, proxy.Listen.Port)
	listen, err := net.Listen("tcp", endpoint)
	if err != nil {
		slog.Error(fmt.Sprintf("【%s】监听 %s 失败！", proxyName, endpoint))
		return
	}
	// 与目标服务器建立连接
	targetEndpoint := fmt.Sprintf("%s:%d", proxy.Target.Host, proxy.Target.Port)
	proxyConn, err := net.Dial("tcp", targetEndpoint)
	if err != nil {
		slog.Error(fmt.Sprintf("【%s】连接代理服务器[%s]失败，错误：%s\n！", proxyName, targetEndpoint, err))
		return
	}
	defer proxyConn.Close()
	slog.Info(fmt.Sprintf("【%s】代理成功，正在监听：%s！", proxyName, endpoint))
	for {
		// 处理客户端连接
		conn, err := listen.Accept()
		if err != nil {
			slog.Error(fmt.Sprintf("【%s】处理客户端连接失败：%s", proxyName, err))
			continue
		}
		go handler(proxyConn, conn, proxyName, proxy)
	}
}

func handler(proxyConn, conn net.Conn, proxyName string, proxy *Proxy) {
	slog.Info(fmt.Sprintf("【%s】与客户端连接成功：%s", proxyName, conn.RemoteAddr().String()))
	defer conn.Close()
	var data = make([]byte, proxy.ReadSize)
	for {
		n, err := conn.Read(data)
		if err != nil {
			if err == io.EOF {
				// 数据读完了
				continue
			}
			// 断开连接
			slog.Error(fmt.Sprintf("【%s】System Error：%s", proxyName, err))
			return
		}
		// 向代理服务器写入数据
		write, err := proxyConn.Write(data[:n])
		if err != nil {
			slog.Error(fmt.Sprintf("【%s】Write Data Error %s", proxyName, err))
			continue
		}
		slog.Info(fmt.Sprintf("【%s】Write Data Success, Byte Num %d！", proxyName, write))
		go func() {
			var data2 = make([]byte, proxy.ReadSize)
			// 将目标服务器写入的数据，copy到客户端。
			n, err = proxyConn.Read(data2)
			if err != nil {
				slog.Error(fmt.Sprintf("【%s】Read Proxy Data Error %s", proxyName, err))
				return
			}
			// 将数据写给客户端
			conn.Write(data2[:n])
		}()
	}
}
